###############################################################################
## AUTHOR: SHOM MAZUMDER
## DATE LAST UPDATED: 07/31/2020
## PURPOSE: MEDIATION EFFECTS
###############################################################################
rm(list = ls())

#### todo ####


#### PACKAGES ####
library(pacman)
p_load(tidyverse,
       estimatr,
       broom,
       haven,
       hrbrthemes,
       cjoint,
       formula.tools, 
       DeclareDesign,
       emmeans,
       lmtest,
       cregg,
       gridExtra,
       egg,
       car,
       DirectEffects,
       purrr,
       modelr,
       rsq)

#### FUNCTIONS ####

# ggplot theme
theme_shom_alt <- function (base_size = 11, waffle = FALSE) 
{
  ret <- theme_minimal(base_size = base_size, base_family = "Roboto Condensed") + 
    theme(plot.background = element_rect(fill = "#f5f5f2", 
                                         color = NA), 
          panel.background = element_rect(fill = "#f5f5f2",color = NA), 
          legend.background = element_rect(fill = "#f5f5f2",color = NA))
  if (waffle) {
    ret + theme(axis.text = element_blank())
  }
  else {
    ret
  }
}

#### LOAD DATA ####
df <- read_rds("01-yougov-conjoint/01-data/cleaned/yougov-conjoint-cleaned.rds")

#### PREP DATA ####

#filter to contests between baseline and co-determinination
df %>%
  filter(cje_corporate_gov %in% c("Privately owned by non-worker shareholders",
                                  "Workers sit on the corporate board")) %>%
  mutate(treated = ifelse(cje_corporate_gov == "Workers sit on the corporate board", 1, 0)) -> df_codetermination

#filter to contests between baseline and esops
df %>%
  filter(cje_corporate_gov %in% c("Privately owned by non-worker shareholders",
                                  "Workers are shareholders")) %>%
  mutate(treated = ifelse(cje_corporate_gov == "Workers are shareholders", 1, 0)) -> df_esop

#filter to contests between baseline and manager elections
df %>%
  filter(cje_corporate_gov %in% c("Privately owned by non-worker shareholders",
                                  "Workers elect their managers")) %>%
  mutate(treated = ifelse(cje_corporate_gov == "Workers elect their managers", 1, 0)) -> df_mgmtelect

#### MEDIATION ANALYSES (CODETERMINATION) ####
set.seed(10009)
num.bootstraps <- 1000

#check mediator interaction assumption
codetermation_tm_inter <- lm(work_binary ~ treated*power_binary,
                                    data = df_codetermination)
summary(codetermation_tm_inter)

#mediation codetermination
direct_codetermination <- sequential_g(work_binary ~ treated | cje_income + cje_firm_size + 
                                         cje_healthcare + cje_hours + cje_retirement + 
                                         cje_work_culture| power_binary, 
                                       data = df_codetermination)
summary(direct_codetermination)

#sensitivity analysis
direct_codetermination_sens <- cdesens(direct_codetermination, var = "treated")
plot(direct_codetermination_sens)

#bootstrap to get acde and anie
df_codetermination %>% 
  modelr::bootstrap(n = num.bootstraps, id = 'boot_num') %>%
  group_by(boot_num) %>%
  mutate(fit_acde = map(strap, ~sequential_g(work_binary ~ treated | cje_income + cje_firm_size + 
                                               cje_healthcare + cje_hours + cje_retirement + 
                                               cje_work_culture| power_binary, 
                                             data = .))) %>%
  mutate(fit_acme = map(strap, ~lm_robust(work_binary ~ treated, data = .))) %>%
  mutate(acde_boot = map(fit_acde, ~coef(.)["treated"]),
         acme_boot = map(fit_acme, ~coef(.)["treated"])) %>%
  select(boot_num,acde_boot,acme_boot) %>%
  unnest %>%
  mutate(anie_boot = acme_boot - acde_boot) %>%
  mutate(treatment = "Workers sit on the corporate board") -> boot_codetermination

#reshape
boot_codetermination_long <- boot_codetermination %>%
  gather(key = "estimand",value = "estimate",-treatment,-boot_num)

#### MEDIATION ANALYSES (ESOP) ####

#check mediator interaction assumption
esop_tm_inter <- lm(work_binary ~ treated*power_binary,
                    data = df_esop)
summary(esop_tm_inter)

#mediation codetermination
direct_esop <- sequential_g(work_binary ~ treated | cje_income + cje_firm_size + 
                                         cje_healthcare + cje_hours + cje_retirement + 
                                         cje_work_culture| power_binary, 
                                       data = df_esop)
summary(direct_esop)

#sensitivity analysis
direct_esop_sens <- cdesens(direct_esop, var = "treated")
plot(direct_esop_sens)

#bootstrap to get acde and anie
df_esop %>% 
  modelr::bootstrap(n = num.bootstraps, id = 'boot_num') %>%
  group_by(boot_num) %>%
  mutate(fit_acde = map(strap, ~sequential_g(work_binary ~ treated | cje_income + cje_firm_size + 
                                               cje_healthcare + cje_hours + cje_retirement + 
                                               cje_work_culture| power_binary, 
                                             data = .))) %>%
  mutate(fit_acme = map(strap, ~lm_robust(work_binary ~ treated, data = .))) %>%
  mutate(acde_boot = map(fit_acde, ~coef(.)["treated"]),
         acme_boot = map(fit_acme, ~coef(.)["treated"])) %>%
  select(boot_num,acde_boot,acme_boot) %>%
  unnest %>%
  mutate(anie_boot = acme_boot - acde_boot) %>%
  mutate(treatment = "Workers are shareholders") -> boot_esop

#reshape
boot_esop_long <- boot_esop %>%
  gather(key = "estimand",value = "estimate",-treatment,-boot_num)

#### MEDIATION ANALYSES (MGMT ELEC) ####

#check mediator interaction assumption
mgmtelect_tm_inter <- lm(work_binary ~ treated*power_binary,
                           data = df_mgmtelect)
summary(mgmtelect_tm_inter)

#mediation codetermination
direct_mgmtelect <- sequential_g(work_binary ~ treated | cje_income + cje_firm_size + 
                              cje_healthcare + cje_hours + cje_retirement + 
                              cje_work_culture| power_binary, 
                            data = df_mgmtelect)
summary(direct_mgmtelect)

#sensitivity analysis
direct_mgmtelect_sens <- cdesens(direct_mgmtelect, var = "treated")
plot(direct_mgmtelect_sens)

#bootstrap to get acde and anie
df_mgmtelect %>% 
  modelr::bootstrap(n = num.bootstraps, id = 'boot_num') %>%
  group_by(boot_num) %>%
  mutate(fit_acde = map(strap, ~sequential_g(work_binary ~ treated | cje_income + cje_firm_size + 
                                               cje_healthcare + cje_hours + cje_retirement + 
                                               cje_work_culture| power_binary, 
                                             data = .))) %>%
  mutate(fit_acme = map(strap, ~lm_robust(work_binary ~ treated, data = .))) %>%
  mutate(acde_boot = map(fit_acde, ~coef(.)["treated"]),
         acme_boot = map(fit_acme, ~coef(.)["treated"])) %>%
  select(boot_num,acde_boot,acme_boot) %>%
  unnest %>%
  mutate(anie_boot = acme_boot - acde_boot) %>%
  mutate(treatment = "Workers elect their managers") -> boot_mgmtelect

#reshape
boot_mgmtelect_long <- boot_mgmtelect %>%
  gather(key = "estimand",value = "estimate",-treatment,-boot_num)



#### GENERATE INTERACTION RESULTS (TABLE S2) ####
col_labels = c("Codetermination",
               "ESOP",
               "Management elections")
dep_var_labels <- c("Work")
covariate_labels = c("Workplace democracy proposal",
                     "Answer to power question",
                     "Workplace democracy proposal X Answer to power question")

stargazer::stargazer(
  codetermation_tm_inter,
  esop_tm_inter,
  mgmtelect_tm_inter,
  column.labels = col_labels,
  covariate.labels = covariate_labels,
  dep.var.labels = dep_var_labels,
  title = "Assessing the interaction effect assumption for causal mediation estimation",
  label = "tab:interaction-test-mediation-assumption"
)

#### GENERATE DISTRIBUTIONS OF QUANTITIES OF INTEREST (FIGURE 4) ####

boot_tbl <- bind_rows(boot_codetermination_long,
                      boot_esop_long,
                      boot_mgmtelect_long) %>%
  mutate(estimand_clean = case_when(
    estimand == "acde_boot" ~ "ACDE",
    estimand == "acme_boot" ~ "AMCE",
    estimand == "anie_boot" ~ "ANIE"
  )) %>%
  mutate(treatment = fct_relevel(treatment,
                                 "Workers sit on the corporate board",
                                 "Workers are shareholders",
                                 "Workers elect their managers"))

boot_tbl %>%
  ggplot(aes(x = estimate,group = estimand_clean,fill = estimand_clean,
             color = estimand_clean)) + 
  geom_vline(xintercept = 0,linetype = "dashed",size = 1.5) + 
  geom_density(alpha = 0.5) + 
  facet_wrap(~treatment) + 
  theme_shom_alt(base_size = 16) + 
  theme(plot.caption = element_text(size = 10,hjust = 1),
        plot.caption.position = "plot",
        strip.background = element_rect(fill="grey"),
        strip.text = element_text(color = 'black'),
        legend.position = "bottom",
        strip.text.x = element_text(size = 11),
        strip.text.y = element_text(size = 11)) + 
  labs(x = 'Estimate',
       y = "density",
       fill = "Estimand",
       color = "Estimand",
       caption = "Notes: This figure plots the distribution of estimated AMCE, ACDE, and ANIEs via the non-parametric bootstrap with 1000 draws.
       We estimate the AMCE via Ordinary Lease Squares compared to the baseline condition of private ownership by non-worker shareholders.
       We estimate the ACDE via sequential-g estimation via the DirectEffects R package.
       The estimate of the ANIE comes from subtracting the bootstrapped ACDE from the bootstrapped AMCE.") + 
  ggsave("01-yougov-conjoint/03-src/output/figures/mediation-effects.pdf",
         dpi = 320,width = 10,height = 6,device = cairo_pdf)

#### PLOT SENSITIVITY ANALYSES (FIGURE S5) ####

# make dataset
codeterm_acde_sens_tbl <- tibble(
    rho = direct_codetermination_sens$rho,
    acde = direct_codetermination_sens$acde,
    se = direct_codetermination_sens$se,
    treatment = "Workers sit on the corporate board"
)

esop_acde_sens_tbl <- tibble(
  rho = direct_esop_sens$rho,
  acde = direct_esop_sens$acde,
  se = direct_esop_sens$se,
  treatment = "Workers are shareholders"
)

mgmtelect_acde_sens_tbl <- tibble(
  rho = direct_mgmtelect_sens$rho,
  acde = direct_mgmtelect_sens$acde,
  se = direct_mgmtelect_sens$se,
  treatment = "Workers elect their managers"
)

sens_tbl <- bind_rows(codeterm_acde_sens_tbl,
                      esop_acde_sens_tbl,
                      mgmtelect_acde_sens_tbl) %>%
  mutate(treatment = fct_relevel(treatment,
                                 "Workers sit on the corporate board",
                                 "Workers are shareholders",
                                 "Workers elect their managers"))

rsq_codeterm_model <- lm_robust(work_binary ~ cje_income + cje_firm_size + 
                    cje_healthcare + cje_hours + cje_retirement + 
            cje_union + cje_retirement + cje_sick_leave,
        data = df_codetermination) %>%
  rsquare(.,df_codetermination)

rsq_mgmtelect_model <- lm_robust(work_binary ~ cje_income + cje_firm_size + 
                                  cje_healthcare + cje_hours + cje_retirement + 
                                  cje_union + cje_retirement + cje_sick_leave,
                                data = df_mgmtelect) %>%
  rsquare(.,df_mgmtelect)

rsq_esop_model <- lm_robust(work_binary ~ cje_income + cje_firm_size + 
                                  cje_healthcare + cje_hours + cje_retirement + 
                                  cje_union + cje_retirement + cje_sick_leave,
                                data = df_esop) %>%
  rsquare(.,df_esop)

rho_tbl <- tibble(
  rho = sqrt(c(rsq_codeterm_model,rsq_esop_model,rsq_mgmtelect_model)),
  treatment = c("Workers sit on the corporate board",
                "Workers are shareholders",
                "Workers elect their managers")
) %>%
  mutate(treatment = fct_relevel(treatment,
                                 "Workers sit on the corporate board",
                                 "Workers are shareholders",
                                 "Workers elect their managers"))

sens_tbl %>%
  ggplot(aes(x = rho,y = acde)) +
  geom_ribbon(aes(ymin = acde - 1.96*se,
                ymax = acde + 1.96*se),
              fill = "indianred",alpha = 0.3) + 
  geom_ribbon(aes(ymin = acde - 1.64*se,
                  ymax = acde + 1.64*se),
              fill = "indianred",alpha = 0.7) + 
  geom_hline(yintercept = 0,linetype = 'dashed',size = 1.5) + 
  geom_vline(xintercept = 0,linetype = 'dashed',size = 1.5) + 
  geom_line(size = 1) +
  theme_shom_alt(base_size = 16) + 
  theme(plot.caption = element_text(size = 10,hjust = 1),
        plot.caption.position = "plot",
        strip.background = element_rect(fill="grey"),
        strip.text = element_text(color = 'black'),
        legend.position = "bottom",
        strip.text.x = element_text(size = 11),
        strip.text.y = element_text(size = 11)) + 
  geom_point(data = rho_tbl,
             aes(x = rho,y = 0),
             shape = 8,size = 3) + 
  geom_point(data = rho_tbl,
             aes(x = -rho,y = 0),
             shape = 8,size = 3) + 
  facet_wrap(~treatment) +
  labs(x = "Correlation between Mediator and Outcome Errors",
       y = "Estimated ACDE",
       caption = "Notes: The dark shaded region represents the 90% confidence interval while the light shaded region represents the 95% confidence interval.
       Points represent the total correlation of the predictive power of economic treatments and the outcome.") + 
  ggsave("01-yougov-conjoint/03-src/output/figures/mediation-sensitivity-analysis.pdf",
         dpi = 320,width = 10,height = 6,device = cairo_pdf)

